#!/usr/local/bin/python

"""
Compiles scripts for running rseg on lustre Submits to lustre 
Input is bamfiles, output will go to dir current dir
"""

import os
import subprocess
import sys
import re

# ----------------------------[ Settings ]--------------------------------------

BAMDIR= '../chipseq_pipeline/bam_clean'             ## Dir with bamfiles
CTRLFILE= 'bham385.bwa.mm9.bed'                  ## The control file to use. Assumed to be in *BAMDIR*. Note that it must be a *BED* file
SKIPBAMS= ['bham385.bwa.mm9.bam'] ## List of bamfiles in BAMDIR to skip (e.g. the control file)
OUTDIR= 'output_mode1_bin-default' ## Where output, including sh scripts and log files will be
JOBNAME= 'rseg4-' ## Prefix for the job names (passed to bsub -J)
# ------------------------------------------------------------------------------

p= subprocess.Popen('mkdir -p %s' %(OUTDIR), shell= True)
p.wait()

## p= subprocess.Popen('mkdir -p rseg') ## Where output files will be
bamfiles= sorted(os.listdir(BAMDIR))
bamfiles= [x for x in bamfiles if re.search('\.bam$', x)]
ctrlfile_path= os.path.join(BAMDIR, CTRLFILE)
if not os.path.exists(ctrlfile_path):
    sys.exit('%s: Cannot find control file: %s' %(sys.argv[0], ctrlfile_path)) 

n= 1
for b in bamfiles:
    if b in SKIPBAMS:
        continue
    bam= os.path.join(BAMDIR, b)
    bedfile= os.path.join(OUTDIR, re.sub('\.bam$', '', b) + '.bed')
    shfilename= os.path.join(OUTDIR, re.sub('\.bam$', '', b) + '.sh')
    shfile= open(shfilename, 'w')
    cmd= """
#!/bin/sh

set -e

bamToBed -i %(bam)s > %(bedfile)s

## /home/berald01/applications/rseg/rseg-bin/bin/rseg-diff -v -d /lustre/sblab/berald01/Tritume/deadzones-k40-mm9.bed -c /home/berald01/applications/rseg/auxiliary_files/mouse-mm9-size.bed -mode 2 %(bedfile)s %(ctrlfile_path)s
/home/berald01/applications/rseg/rseg-bin/bin/rseg -v -d /lustre/sblab/berald01/Tritume/deadzones-k40-mm9.bed -c /home/berald01/applications/rseg/auxiliary_files/mouse-mm9-size.bed %(bedfile)s

rm %(bedfile)s
    """ %{'bam':bam, 'bedfile':bedfile, 'ctrlfile_path':ctrlfile_path}
    shfile.write(cmd)
    shfile.close()
    os.chmod(shfilename, 0744)
    jobid= JOBNAME + str(n)
    cmd= """ bsub -R "rusage[mem=8192]" -w deadz -o %(jobid)s.log -J %(jobid)s < %(shfilename)s""" %{'jobid':jobid, 'shfilename':shfilename}
    print(cmd)
    p= subprocess.Popen(cmd, shell= True)
    n += 1
sys.exit()